import oneflow as flow
import oneflow.nn as nn
import oneflow.nn.functional as F


class ConformerConvolutionModule(nn.Module):
    def __init__(self, channels, kernel_size, bias=True, dropout=0.0):
        super(ConformerConvolutionModule, self).__init__()

        self.pointwise_conv1 = nn.Linear(channels, 2 * channels, bias=bias)

        self.depthwise_conv = nn.Conv1d(
            channels,
            channels,
            kernel_size,
            stride=1,
            padding=(kernel_size - 1) // 2,
            groups=channels,
            bias=bias
        )

        self.batch_norm = nn.BatchNorm1d(channels)

        self.pointwise_conv2 = nn.Linear(channels, channels, bias=bias)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        """
        Args:
            x: [batch_size, time, channels]
            mask: [batch_size, time]
        """
        mask = mask.unsqueeze(2).repeat([1, 1, x.size(-1)])

        x = self.pointwise_conv1(x)
        x = F.glu(x)
        # x.masked_fill_(~mask, 0.0)
        x = flow.masked_fill(x,mask==0, 0.0)

        x = x.transpose(1, 2)
        x = self.depthwise_conv(x)
        x = self.batch_norm(x)
        x = x * flow.sigmoid(x) # swish
        x = x.transpose(1, 2)

        x = self.pointwise_conv2(x)
        # x.masked_fill_(~mask, 0.0)
        x = flow.masked_fill(x,mask==0, 0.0)

        return x


class UniDirectionalConformerConvolutionModule(nn.Module):
    def __init__(self, channels, kernel_size, bias=True, dropout=0.0):
        super(UniDirectionalConformerConvolutionModule, self).__init__()

        self.kernel_size = kernel_size

        self.pointwise_conv1 = nn.Linear(channels, 2 * channels, bias=bias)

        self.depthwise_conv = nn.Conv1d(
            channels,
            channels,
            kernel_size,
            stride=1,
            padding=0,
            groups=channels,
            bias=bias
        )

        self.batch_norm = nn.BatchNorm1d(channels)

        self.pointwise_conv2 = nn.Linear(channels, channels, bias=True)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        """
        Args:
            x: [batch_size, time, channels]
            mask: [batch_size, time]
        """
        mask = mask.unsqueeze(2).repeat([1, 1, x.size(-1)])
        
        x = self.pointwise_conv1(x)
        x = F.glu(x)
        # x.masked_fill_(~mask, 0.0)
        x = flow.masked_fill(x,mask==0, 0.0)

        x = F.pad(x, pad=(0, 0, self.kernel_size - 1, 0), value=0.0).transpose(1, 2)
        x = self.depthwise_conv(x)
        x = self.batch_norm(x)
        x = x * flow.sigmoid(x) # swish
        x = x.transpose(1, 2)

        x = self.pointwise_conv2(x)
        # x.masked_fill_(~mask, 0.0)
        x = flow.masked_fill(x,mask==0, 0.0)

        return x

    def inference(self):
        raise NotImplementedError